import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import umap

if __name__ == '__main__':
    penguins = pd.read_csv("penguins_size.csv")
    print(penguins.head())
    penguins = penguins.dropna()
    print(penguins.species_short.value_counts())
    # sns.pairplot(penguins, hue='species_short')
    # plt.show()
    reducer = umap.UMAP()
    penguin_data = penguins[
        [
            "culmen_length_mm",
            "culmen_depth_mm",
            "flipper_length_mm",
            "body_mass_g",
        ]
    ].values
    scaled_penguin_data = StandardScaler().fit_transform(penguin_data)
    embedding = reducer.fit_transform(scaled_penguin_data)
    print(embedding)
    plt.scatter(
        embedding[:, 0],
        embedding[:, 1],
        c=[sns.color_palette()[x] for x in penguins.species_short.map({"Adelie": 0, "Chinstrap": 1, "Gentoo": 2})])
    plt.gca().set_aspect('equal', 'datalim')
    plt.title('UMAP projection of the Penguin dataset', fontsize=24)
    plt.show()
